import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


plt.rcParams['figure.figsize'] = (7, 3.5)
plt.rcParams['figure.dpi'] = 150
plt.rcParams['axes.unicode_minus'] = False  #解决坐标轴负数的铅显示问题

"""
leaky_relu(x) = x if x >=0 else \alpha\times(e^x-1)

grad(x)  = 1 if x >= 0 else \alpha\timese^x

理想的激活函数应满足两个条件：
    - 输出的分布是零均值的，可以加快训练速度。
    - 激活函数是单侧饱和的，可以更好的收敛。

LeakyReLU和PReLU满足第1个条件，不满足第2个条件；而ReLU满足第2个条件，不满足第1个条件。两个条件都满足的激活函数为ELU(Exponential 
Linear Unit)。ELU虽然也不是零均值的，但在以0为中心一个较小的范围内，均值是趋向于0，当然也与α \alphaα的取值也是相关的。

- 优点
    - ELU具有Relu的大多数优点，不存在Dead Relu问题，输出的均值也接近为0值；
    - 该函数通过减少偏置偏移的影响，使正常梯度更接近于单位自然梯度，从而使均值向0加速学习；
    - 该函数在负数域存在饱和区域，从而对噪声具有一定的鲁棒性；

- 缺点
    - 计算强度较高，含有幂运算；
    - 在实践中同样没有较Relu更突出的效果，故应用不多；
"""

#! 1. 定义可以进行前项和后算运算LeakyReLU算子
class MyELUAutoGrad(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, alpha):
        ctx.alpha = alpha
        ctx.save_for_backward(x)       # 保存上下文用于反向传播
        return F.elu(input=x, alpha=alpha) 
    
    @staticmethod
    def backward(ctx, grad_output):    # y.backward(torch.ones_like(y)) 传进来的
        x,     = ctx.saved_tensors
        ones_x = torch.ones_like(x)
        grad   = torch.where(x < 0, torch.exp(x) * ctx.alpha, ones_x)
        return grad, None
    

class ELU(nn.Module):
    def __init__(self, alpha):
        super(ELU, self).__init__()
        self.alpha = alpha
    
    def forward(self, x):
        return MyELUAutoGrad.apply(x, self.alpha)
    
    
x   = torch.randn(2, 10, dtype=torch.float32, requires_grad=True)
print(x)

net = ELU(alpha=0.1)
y   = net(x)
y.backward(torch.ones_like(y))       

print(x.grad)
print("=========================")


model = nn.ELU(alpha=0.1)
x.grad.zero_()
y    = model(x)
y.backward(torch.ones_like(y))       

print(x.grad)